﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

using System.Diagnostics.Contracts;
using HeapModel.ProgInfo;
using HeapModel.ConcreteMemory;
using HeapModel.RegionPartition;
using System.Diagnostics;
using System.IO;
using System.Xml;

namespace HeapModel.ReducedGraph
{
	/// <summary>
	/// This class represents a reduced abstract graph.
	/// These are regions that are further abstracted based on dominance and structural equivalence for the storage similarity relation.
	/// </summary>
	[DebuggerDisplay("Node = {NodeCount}")]
	public sealed class ReducedHeapGraph
	{
		public readonly RegionGraph BaseHG;
		
		private readonly HashSet<DomReducedNode> m_nodes;
		private readonly Dictionary<RegionNode, DomReducedNode> m_containmentmap;

		private readonly Dictionary<DomReducedNode, HashSet<DomReducedNode>> m_finalrelmap;

		[ContractInvariantMethod]
		private void ReduceHeapGraphInvariant()
		{
			Contract.Invariant(this.BaseHG != null);
			Contract.Invariant(this.m_nodes != null);
			Contract.Invariant(this.m_containmentmap != null);
			Contract.Invariant(this.m_finalrelmap != null);
		}

		private ReducedHeapGraph(RegionGraph hg)
		{
			Contract.Requires(hg != null);

			this.BaseHG = hg;
			this.m_nodes = new HashSet<DomReducedNode>();
			this.m_containmentmap = new Dictionary<RegionNode, DomReducedNode>();
			this.m_finalrelmap = new Dictionary<DomReducedNode, HashSet<DomReducedNode>>();
		}

		/// <summary>
		/// Create and reduce based on the input region graph.
		/// </summary>
		public static ReducedHeapGraph Create(RegionGraph hg)
		{
			Contract.Requires(hg != null);
			Contract.Ensures(Contract.Result<ReducedHeapGraph>() != null);

			var rhg = InitializeAbstractHeapGraph(hg);
			rhg.CollapseCompleteDominators();
			
			rhg.ConstructFinalRelMap();
			return rhg;
		}

		/// <summary>
		/// Create and initialize a ReducedHeapGraph for the given base HeapGraph.
		/// </summary>
		private static ReducedHeapGraph InitializeAbstractHeapGraph(RegionGraph hg)
		{
			Contract.Requires(hg != null);
			Contract.Ensures(Contract.Result<ReducedHeapGraph>() != null);

			ReducedHeapGraph rdhg = new ReducedHeapGraph(hg);

			//gen all the new nodes (but don't set in/out edges yet)
			foreach(RegionNode hr in hg.Regions)
			{
				DomReducedNode rr = DomReducedNode.InitializeReducedForGivenRegion(hr);
				rdhg.m_containmentmap.Add(hr, rr);
				rdhg.m_nodes.Add(rr);
			}

			return rdhg;
		}

		/// <summary>
		/// Writes the Dgml graph to a file.
		/// </summary>
		public void WriteDgml(String fileName, String basefn)
		{
			Contract.Requires(!String.IsNullOrEmpty(fileName));

			this.BaseHG.WriteDgml(basefn);

			using(var writer = new StreamWriter(fileName))
				this.WriteDgml(writer);
		}

		/// <summary>
		/// Write the graph to a writer as DGML
		/// </summary>
		public void WriteDgml(TextWriter writer)
		{
			Contract.Requires(writer != null);

			int totalheapbytes = this.BaseHG.BaseConcreteHeap.AllObjects.Aggregate<ConcreteNode, int>(0, (int acc, ConcreteNode cn) => acc + cn.GlueSizeInBytes + cn.DataSizeInBytes);
			using(var xw = XmlWriter.Create(writer, new XmlWriterSettings { Indent = true }))
			{
				xw.WriteStartDocument();
				xw.WriteStartElement("DirectedGraph", DgmlExtensions.NamespaceUri);
				xw.WriteStartElement("Nodes");
				this.WriteNodesIntoDgml(xw, totalheapbytes);
				xw.WriteEndElement();

				xw.WriteStartElement("Links");
				this.WriteLinksIntoDgml(xw);
				xw.WriteEndElement(); // Links
				DgmlExtensions.WriteCategoriesAndProperties(xw);
				DgmlExtensions.WriteStyles(xw);
				xw.WriteEndElement(); // DirectedGraph
				xw.WriteEndDocument();
			}
		}

		/// <summary>
		/// Write all the node information in this into the xmlwriter.
		/// </summary>
		public void WriteNodesIntoDgml(XmlWriter xmlw, int totalbytesinheap)
		{
			this.BaseHG.WriteNodesIntoDgml(xmlw);

			foreach(DomReducedNode rr in this.m_nodes)
				rr.WriteDgml(totalbytesinheap, xmlw);
		}

		/// <summary>
		/// Write all the link info from  this into the xmlwriter
		/// </summary>
		public void WriteLinksIntoDgml(XmlWriter xmlwriter)
		{
			this.BaseHG.WriteLinksIntoDgml(xmlwriter);

			foreach(DomReducedNode rr in this.m_nodes)
			{
				if(rr.ContainedRegions.Count > 1)
				{
					foreach(RegionNode hr in rr.ContainedRegions)
					{
						xmlwriter.WriteStartElement("Link");
						xmlwriter.WriteAttributeString("Source", "r" + rr.NodeID.ToString());
						xmlwriter.WriteAttributeString("Target", hr.RegionID.ToString());
						xmlwriter.WriteAttributeString("Category", "Contains");
						xmlwriter.WriteEndElement();
					}
				}
			}
		}

		internal HashSet<DomReducedNode> Nodes
		{ get { return this.m_nodes; } }

		internal RegionGraph BaseHeapGraph
		{ get { return this.BaseHG; } }

		/// <summary>
		/// Return the relation between the nodes based on the cross region pointers.
		/// </summary>
		internal Dictionary<DomReducedNode, HashSet<DomReducedNode>> RegionRelationMap
		{
			get
			{
				Contract.Ensures(Contract.Result<Dictionary<DomReducedNode, HashSet<DomReducedNode>>>() != null);

				return this.m_finalrelmap;
			}
		}

		private void ConstructFinalRelMap()
		{
			foreach(DomReducedNode drn in this.m_nodes)
			{
				HashSet<DomReducedNode> trgts = new HashSet<DomReducedNode>();
				foreach(RegionNode rgn in drn.ContainedRegions)
				{
					foreach(RegionNode trgt in this.BaseHG.RegionRelationMap[rgn])
					{
						DomReducedNode trgtrgn = this.m_containmentmap[trgt];
						if(trgtrgn != drn)
							trgts.Add(trgtrgn);
					}
				}

				this.m_finalrelmap.Add(drn, trgts);
			}
		}

		/////////////////////////////////////////////////////////
		//traversal support

		//time stamp for doing dfs visit
		private int m_ftime;

		/// <summary>
		/// Prep for DFS.
		/// </summary>
		private void PrepForDFS()
		{
			foreach(DomReducedNode nn in this.m_nodes)
				nn.NodeColor = GraphColor.White;
		}

		/// <summary>
		/// Visit a single node in the DFS traversal
		/// </summary>
		private void DoDFSVisit(DomReducedNode n)
		{
			Contract.Requires(n != null);

			if(n.NodeColor != GraphColor.White)
				return;

			n.NodeColor = GraphColor.Gray;
			foreach(RegionNode nn in n.ContainedRegions)
			{
				foreach(RegionNode trgtn in this.BaseHG.RegionRelationMap[nn])
				{
					DomReducedNode drn = this.m_containmentmap[trgtn];
					this.DoDFSVisit(drn);
				}
			}

			n.NodeColor = GraphColor.Black;
			n.FinishTime = this.m_ftime++;
		}

		/// <summary>
		/// Do the DFS of the graph.
		/// </summary>
		private void DoDFS()
		{
			this.PrepForDFS();
			this.m_ftime = 1;
			foreach(DomReducedNode drn in this.m_nodes)
				this.DoDFSVisit(drn);
		}

		/// <summary>
		/// Sort the graph psudo-topologically put the results into sorted.
		/// e.g. the result is topologically sorted outside of any SCC's and inside SCC's some order is selected which is a valid topological order 
		/// if we pick one of the entry points to the as the largest element and then order the SCC starting from this node ignoring any back edges.
		/// </summary>
		private void TopoSort(List<DomReducedNode> sorted)
		{
			Contract.Requires(sorted != null);
			Contract.Requires(sorted.Count == 0);

			this.DoDFS();

			foreach(DomReducedNode nn in this.m_nodes)
				sorted.Add(nn);

			sorted.Sort((DomReducedNode r1, DomReducedNode r2) => r2.FinishTime - r1.FinishTime);
		}

		/////////////////////////////////////////////////////////
		//Normalize Routines
		
		/// <summary>
		/// Find dominators and collapse the dominated set.
		/// </summary>
		private void CollapseCompleteDominators()
		{
			Dictionary<DomReducedNode, HashSet<DomReducedNode>> preddict = new Dictionary<DomReducedNode, HashSet<DomReducedNode>>();
			foreach(DomReducedNode drn in this.m_nodes)
				preddict.Add(drn, new HashSet<DomReducedNode>());

			foreach(DomReducedNode drn in this.m_nodes)
			{
				foreach(RegionNode rn in drn.ContainedRegions)
				{
					foreach(RegionNode trgtn in this.BaseHG.RegionRelationMap[rn])
					{
						DomReducedNode trgtdrn = this.m_containmentmap[trgtn];
						preddict[trgtdrn].Add(drn);
					}
				}
			}

			this.MarkDomSafeReducedNodes(preddict);

			Dictionary<DomReducedNode, HashSet<DomReducedNode>> dominatorMap = this.ComputeAllNodesDomByDict(preddict);

			List<DomReducedNode> tpolist = new List<DomReducedNode>();
			this.TopoSort(tpolist);

			List<DomReducedNode> toproc = new List<DomReducedNode>(tpolist);
			toproc.Reverse(); //then we read from back in worklist below
			while(toproc.Count != 0)
			{
				DomReducedNode crgn = toproc[toproc.Count - 1];
				HashSet<DomReducedNode> doms = new HashSet<DomReducedNode>(dominatorMap[crgn]);

				if((!crgn.IsDomSafe) && doms.Count > 1) // && (doms.Count < tpolist.Count / 2))
				{
					//Remove all the nodes that are going to be merged from the dominator and reachable sets
					toproc.RemoveAll((DomReducedNode r) => doms.Contains(r));

					DomReducedNode nhr = DomReducedNode.JoinNodes(doms, crgn.RootRgnTypes);

					foreach(RegionNode rn in nhr.ContainedRegions)
						this.m_containmentmap[rn] = nhr;

					this.m_nodes.ExceptWith(doms);
					this.m_nodes.Add(nhr);

					UpdateReachableOrDomMapAfterJoin(doms, dominatorMap);
				}
				else
					toproc.RemoveAt(toproc.Count - 1);
			}
		}

		/// <summary>
		/// A helper for above method to update reachable and dom sets after replacing a set of nodes with a new summary node.
		/// Do not let replset alias with an entry in the dictionary.
		/// </summary>
		private static void UpdateReachableOrDomMapAfterJoin(HashSet<DomReducedNode> replset, Dictionary<DomReducedNode, HashSet<DomReducedNode>> dict)
		{
			Contract.Requires(replset != null);
			Contract.Requires(replset.Count != 0);
			Contract.Requires(dict != null);

			HashSet<DomReducedNode> newrgnentry = new HashSet<DomReducedNode>();

			//compute value for rgn entry and remove all entries for values in replset
			foreach(DomReducedNode ohr in replset)
			{
				newrgnentry.UnionWith(dict[ohr]);
				dict.Remove(ohr);
			}

			//remove all values in replset from the ranges in the map and if they were in the set add newrgn
			foreach(HashSet<DomReducedNode> rrs in dict.Values)
			{
				if(rrs.Overlaps(replset))
					rrs.ExceptWith(replset);
			}
		}

		/// <summary>
		/// Compute a dictionary that maps from ReducedRegions to the set of ReducedRegions that it immediately dominataes (and self).
		/// </summary>
		private Dictionary<DomReducedNode, HashSet<DomReducedNode>> ComputeAllNodesDomByDict(Dictionary<DomReducedNode, HashSet<DomReducedNode>> preddict)
		{
			Contract.Ensures(Contract.Result<Dictionary<DomReducedNode, HashSet<DomReducedNode>>>() != null);

			Dictionary<DomReducedNode, HashSet<DomReducedNode>> doms = this.ComputeDominatorSets(preddict);

			Dictionary<DomReducedNode, HashSet<DomReducedNode>> resd = new Dictionary<DomReducedNode, HashSet<DomReducedNode>>();
			foreach(DomReducedNode rr in this.m_nodes)
			{
				HashSet<DomReducedNode> rrs = new HashSet<DomReducedNode>();
				rrs.Add(rr);
				resd.Add(rr, rrs);
			}

			foreach(KeyValuePair<DomReducedNode, HashSet<DomReducedNode>> domsp in doms)
			{
				foreach(DomReducedNode rr in domsp.Value)
					resd[rr].Add(domsp.Key);
			}

			return resd;
		}

		/// <summary>
		/// Mark nodes that are safe based on var accesses. Larger radius than for standard view.
		/// </summary>
		private void MarkDomSafeReducedNodes(Dictionary<DomReducedNode, HashSet<DomReducedNode>> preddict)
		{
			foreach(KeyValuePair<Tuple<String, bool>, ConcreteNode> crnkvp in this.BaseHG.BaseConcreteHeap.Roots)
			{
				if(crnkvp.Value.ConcreteNodeType.IsSpecialNullType)
					continue;

				DomReducedNode rr = this.m_containmentmap[this.BaseHG.ContainmentMap[crnkvp.Value]];
				rr.IsDomSafe = true;
				foreach(DomReducedNode drn in preddict[rr])
					drn.IsDomSafe = true;

				foreach(RegionNode rn in rr.ContainedRegions)
				{
					foreach(RegionNode trgtn in this.BaseHG.RegionRelationMap[rn])
					{
						if(this.BaseHG.RegionRelationMap[trgtn].Count < 20)
						{
							DomReducedNode rro = this.m_containmentmap[trgtn];
							rro.IsDomSafe = true;
						}
					}
				}
			}
		}

		/// <summary>
		/// Compute the dominator sets for each node.
		/// </summary>
		private Dictionary<DomReducedNode, HashSet<DomReducedNode>> ComputeDominatorSets(Dictionary<DomReducedNode, HashSet<DomReducedNode>> preddict)
		{
			Contract.Ensures(Contract.Result<Dictionary<DomReducedNode, HashSet<DomReducedNode>>>() != null);

			//
			//Impl. taken from Advanced Compiler Design and Implelentation p. 182
			//
			Dictionary<DomReducedNode, HashSet<DomReducedNode>> domin = new Dictionary<DomReducedNode, HashSet<DomReducedNode>>();
			HashSet<DomReducedNode> N = new HashSet<DomReducedNode>(this.m_nodes);

			foreach(DomReducedNode rr in this.m_nodes)
			{
				HashSet<DomReducedNode> doms = new HashSet<DomReducedNode>();
				if(rr.IsDomSafe)
					doms.Add(rr);
				else
					doms.UnionWith(N);
				domin.Add(rr, doms);
			}

			bool change = true;
			while(change)
			{
				change = false;
				foreach(DomReducedNode rr in this.m_nodes)
				{
					if(rr.IsDomSafe)
						continue;

					HashSet<DomReducedNode> T = new HashSet<DomReducedNode>(N);
					foreach(DomReducedNode predrgn in preddict[rr])
						T.IntersectWith(domin[predrgn]);

					HashSet<DomReducedNode> D = T;
					D.Add(rr);
					if(!domin[rr].SetEquals(D))
					{
						domin[rr] = D;
						change = true;
					}
				}
			}

			return domin;
		}
	}
}
